import os
import cv2
import numpy as np
from sklearn import svm
from sklearn.model_selection import train_test_split
from char_recognition import preprocess_char, CharRecognizer


def load_data(data_dir):
    X = []
    y = []
    for char in os.listdir(data_dir):
        char_dir = os.path.join(data_dir, char)
        if os.path.isdir(char_dir):
            for img_file in os.listdir(char_dir):
                img_path = os.path.join(char_dir, img_file)
                img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
                if img is not None:
                    features = preprocess_char(img)
                    X.append(features)
                    y.append(char)
    return np.array(X), np.array(y)


def train_model():
    data_dir = '../data/train'
    X, y = load_data(data_dir)

    if len(X) == 0:
        print("No training data found. Please check your data directory.")
        return

    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    recognizer = CharRecognizer()
    recognizer.train(X_train, y_train)

    # 评估模型
    accuracy = recognizer.model.score(X_test, y_test)
    print(f"Model accuracy: {accuracy:.2f}")

    # 保存模型
    model_dir = '../models'
    if not os.path.exists(model_dir):
        os.makedirs(model_dir)
    model_path = os.path.join(model_dir, 'svm_model.dat')
    recognizer.save_model(model_path)
    print(f"Model saved to {model_path}")


if __name__ == "__main__":
    train_model()